{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flocks, Herds, and Traffic Jams\n",
    "\n",
    "Code examples from [Think Complexity, 2nd edition](http://greenteapress.com/wp/complexity2), Chapter 10\n",
    "\n",
    "Copyright 2016 Allen Downey, [MIT License](http://opensource.org/licenses/MIT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import thinkplot\n",
    "from thinkstats2 import RandomSeed\n",
    "\n",
    "from Cell2D import Cell2DViewer\n",
    "\n",
    "from matplotlib import rc\n",
    "rc('animation', html='html5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Driver:\n",
    "    \n",
    "    def __init__(self, loc, speed=4):\n",
    "        \"\"\"Initialize the attributes.\n",
    "        \n",
    "        loc: position on track, in miles\n",
    "        speed: speed in miles per hour\n",
    "        \"\"\"\n",
    "        self.start = loc\n",
    "        self.loc = loc\n",
    "        self.speed = speed\n",
    "        \n",
    "    def choose_acceleration(self, d):\n",
    "        return 1\n",
    "        \n",
    "    def set_odometer(self):\n",
    "        self.start = self.loc\n",
    "        \n",
    "    def read_odometer(self):\n",
    "        return self.loc - self.start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Highway:\n",
    "    \n",
    "    max_acc = 1\n",
    "    min_acc = -10\n",
    "    speed_limit = 40\n",
    "    \n",
    "    def __init__(self, n=10, length=1000, eps=0, constructor=Driver):\n",
    "        \"\"\"Initializes the attributes.\n",
    "        \n",
    "        n: number of drivers\n",
    "        length: length of the track\n",
    "        eps: variability in speed\n",
    "        constructor: function used to instantiate drivers\n",
    "        \"\"\"\n",
    "        self.length = length\n",
    "        self.eps = eps\n",
    "        self.crashes = 0\n",
    "\n",
    "        # create the drivers\n",
    "        locs = np.linspace(0, length, n, endpoint=False)\n",
    "        self.drivers = [constructor(loc) for loc in locs]\n",
    "        \n",
    "        # and link them up\n",
    "        for i in range(n):\n",
    "            j = (i+1) % n\n",
    "            self.drivers[i].next = self.drivers[j]\n",
    "            \n",
    "    def step(self):\n",
    "        \"\"\"Performs one time step.\"\"\"\n",
    "        for driver in self.drivers:\n",
    "            self.move(driver)\n",
    "            \n",
    "    def move(self, driver):\n",
    "        \"\"\"Updates `driver`.\n",
    "        \n",
    "        driver: Driver object\n",
    "        \"\"\"\n",
    "        # get the distance to the next driver\n",
    "        d = self.distance(driver)\n",
    "\n",
    "        # let the driver choose acceleration\n",
    "        acc = driver.choose_acceleration(d)\n",
    "        acc = min(acc, self.max_acc)\n",
    "        acc = max(acc, self.min_acc)\n",
    "        speed = driver.speed + acc\n",
    "            \n",
    "        # add random noise to speed\n",
    "        speed *= np.random.uniform(1-self.eps, 1+self.eps)\n",
    "        \n",
    "        # keep it nonnegative and under the speed limit\n",
    "        speed = max(speed, 0)\n",
    "        speed = min(speed, self.speed_limit)\n",
    "        \n",
    "        # if current speed would collide with next driver, stop\n",
    "        if speed > d:\n",
    "            speed = 0\n",
    "            self.crashes += 1\n",
    "            \n",
    "        # update speed and loc\n",
    "        driver.speed = speed\n",
    "        driver.loc += speed\n",
    "            \n",
    "    def distance(self, driver):\n",
    "        \"\"\"Distance from `driver` to next driver.\n",
    "        \n",
    "        driver: Driver object\n",
    "        \"\"\"\n",
    "        d = driver.next.loc - driver.loc\n",
    "        # fix wraparound\n",
    "        if d < 0:\n",
    "            d += self.length\n",
    "        return d\n",
    "    \n",
    "    def set_odometers(self):\n",
    "        return [driver.set_odometer()\n",
    "                for driver in self.drivers] \n",
    "    \n",
    "    def read_odometers(self):\n",
    "        return np.mean([driver.read_odometer()\n",
    "                        for driver in self.drivers])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HighwayViewer(Cell2DViewer):\n",
    "    \"\"\"Generates visualization and animation of SugarScape.\"\"\"\n",
    "    \n",
    "    def draw(self, grid=False):\n",
    "        \"\"\"Draws the array and any other elements.\n",
    "        \n",
    "        grid: boolean, whether to draw grid lines\n",
    "        \"\"\"\n",
    "        self.draw_drivers()\n",
    "        \n",
    "        plt.axis('off')\n",
    "        plt.axis('equal')\n",
    "        plt.xlim([-1.05, 1.05])\n",
    "        plt.ylim([-1.05, 1.05])\n",
    "\n",
    "    def draw_drivers(self):\n",
    "        \"\"\"Plots the drivers.\n",
    "        \"\"\"\n",
    "        drivers = self.viewee.drivers\n",
    "        xs, ys = self.get_coords(drivers)\n",
    "        self.points = plt.plot(xs, ys, '.', color='blue', markersize=20)[0]\n",
    "        \n",
    "        stopped = [driver for driver in self.viewee.drivers \n",
    "                  if driver.speed==0]\n",
    "        xs, ys = self.get_coords(stopped)\n",
    "        xs *= 0.9\n",
    "        ys *= 0.9\n",
    "        self.stopped = plt.plot(xs, ys, 'x', color='red')[0]\n",
    "        \n",
    "    def animate_func(self, i):\n",
    "        \"\"\"Draws one frame of the animation.\"\"\"\n",
    "        if i > 0:\n",
    "            self.step()\n",
    "        drivers = self.viewee.drivers\n",
    "        xs, ys = self.get_coords(drivers)\n",
    "        self.points.set_data(np.array([xs, ys]))\n",
    "        \n",
    "        stopped = [driver for driver in self.viewee.drivers \n",
    "                  if driver.speed==0]\n",
    "        xs, ys = self.get_coords(stopped)\n",
    "        self.stopped.set_data(np.array([xs, ys]) * 0.9)\n",
    "        return self.points,\n",
    "    \n",
    "    def get_coords(self, drivers):\n",
    "        \"\"\"Gets the coordinates of the drivers.\n",
    "        \n",
    "        Transforms from (row, col) to (x, y).\n",
    "        \n",
    "        returns: tuple of sequences, (xs, ys)\n",
    "        \"\"\"\n",
    "        scale = 2 * np.pi / self.viewee.length\n",
    "        \n",
    "        locs = np.array([driver.loc for driver in drivers])\n",
    "        xs = np.cos(locs * scale)\n",
    "        ys = np.sin(locs * scale)\n",
    "        return xs, ys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "highway = Highway(30, eps=0.02)\n",
    "viewer = HighwayViewer(highway)\n",
    "anim = viewer.animate(frames=600, interval=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(22)\n",
    "\n",
    "highway = Highway(30, eps=0.02)\n",
    "viewer = HighwayViewer(highway)\n",
    "\n",
    "thinkplot.preplot(cols=3)\n",
    "for i in range(16):\n",
    "    highway.step()\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(2)\n",
    "for i in range(1):\n",
    "    highway.step()\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(3)\n",
    "for i in range(1):\n",
    "    highway.step()\n",
    "viewer.draw()\n",
    "\n",
    "plt.savefig('chap10-1.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(eps, constructor=Driver, iters=100):\n",
    "    res = []\n",
    "    for n in range(5, 100, 5):\n",
    "        highway = Highway(n, eps=eps, constructor=constructor)\n",
    "        for i in range(iters):\n",
    "            highway.step()\n",
    "\n",
    "        highway.set_odometers()\n",
    "        for i in range(iters):\n",
    "            highway.step()\n",
    "\n",
    "        res.append((n, highway.read_odometers() / iters))\n",
    "    \n",
    "    return np.transpose(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(20)\n",
    "\n",
    "for eps in [0.0, 0.001, 0.01]:\n",
    "    xs, ys = run_simulation(eps)\n",
    "    thinkplot.plot(xs, ys, label='eps=%g' % eps)\n",
    "    \n",
    "thinkplot.config(xlabel='Number of cars',\n",
    "                 ylabel='Average speed',\n",
    "                 xlim=[0, 100], ylim=[0, 42])\n",
    "\n",
    "plt.savefig('chap10-2.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** In the traffic jam simulation, define a class, `BetterDriver`,\n",
    "that inherits from `Driver` and overrides `choose_acceleration`.\n",
    "See if you can define driving rules that do better than the basic\n",
    "implementation in `Driver`.  You might try to achieve higher\n",
    "average speed, or a lower number of collisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Not actually a solution yet; just playing around.\n",
    "\n",
    "class BetterDriver(Driver):\n",
    "    def choose_acceleration(self, d):\n",
    "        if self.speed < 20:\n",
    "            return 1\n",
    "        else:\n",
    "            return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for constructor in [Driver, BetterDriver]:\n",
    "    xs, ys = run_simulation(eps=0.0, constructor=constructor)\n",
    "    thinkplot.plot(xs, ys, label=constructor.__name__)\n",
    "    \n",
    "thinkplot.config(xlabel='Number of cars',\n",
    "                 ylabel='Average speed',\n",
    "                 xlim=[0, 100], ylim=[0, 42])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
